import luojianet.dataset as ds
#创建cifar数据集对象，需要自己下载解压bin格式的数据集文件
#下载地址：https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/cifar-10-binary.tar.gz
#CIFAR-10数据集文件的目录结构如下：
#datasets/
#├── cifar-10-batches-bin
#│   ├── batches.meta
#│   ├── data_batch_1
#│   ├── data_batch_2
#│   ├── data_batch_3
#│   ├── data_batch_4
#│   ├── data_batch_5
#│   ├── readme.html
#│   └── test_batch
data_dir = "./datasets/cifar-10-batches-bin"
datasets_cifar10 = ds.Cifar10Dataset(dataset_dir=data_dir, usage='train', shuffle=False).batch(4) #定义batch为4的数据集对象
batch_size = datasets_cifar10.get_batch_size()
# 数据迭代器使用
data = next(datasets_cifar10.create_dict_iterator()) # dict形式
print(f"Data type:{type(data['image'])}\nImage shape: {data['image'].shape}, Label: {data['label']}")
data1 = next(datasets_cifar10.create_tuple_iterator()) # tuple形式


# 数据处理及增强
# 数据处理
import numpy as np
import matplotlib.pyplot as plt
import luojianet.dataset.vision.c_transforms as transforms
images = data['image'].asnumpy()
labels = data['label'].asnumpy()

# 展示原图
plt.figure()
for i in range(1,5):
    plt.subplot(2, 2, i)
    plt.title(f"{labels[i-1]}")
    plt.imshow(images[i-1])
plt.show()

# 展示使用numpy进行旋转的图
plt.figure()
for i in range(1, 5):
    plt.subplot(2, 2, i)
    image_trans = np.transpose(images[i-1], (1, 0, 2))
    plt.title(f"{labels[i-1]}")
    plt.imshow(image_trans, interpolation="None")
plt.show()

# 数据增强，使用map操作，展示随机剪裁和水平反转后的图
trans = [
    transforms.RandomCrop((36, 36), (4, 4, 4, 4)),
    transforms.RandomHorizontalFlip(prob=0.5),
]

dataset = ds.Cifar10Dataset(dataset_dir=data_dir, shuffle=False, usage='train')
dataset = dataset.map(operations=trans, input_columns=['image']).batch(4)
data = next(dataset.create_dict_iterator())
images = data['image'].asnumpy()
labels = data['label'].asnumpy()
plt.figure()
for i in range(1,5):
    plt.subplot(2,2,i)
    plt.title(f"{labels[i-1]}")
    plt.imshow(images[i-1])
plt.show()